'''
初始化BaseInfo对象和BaseSeq2Vec对应的DNA与RNA版本
'''
import pandas as pd
from models import BaseInfo, BaseSeq2Vec

OUTPUT_PATH = './user_data/'

# 初始化碱基信息
base_info = BaseInfo()
base_info.save(f'{OUTPUT_PATH}base_info.m')

# 获取碱基序列
mirna_seqdf = pd.read_csv('./datasets/train_dataset/mirna_seq.csv')
# 清除换行符
mirna_seqdf['seq'] = mirna_seqdf['seq'].apply(lambda x:x.strip())
rna_seq_list = mirna_seqdf['seq'].tolist()
dna_seq_list = pd.read_csv('./datasets/train_dataset/gene_seq.csv')['sequence'].tolist()

# 初始化序列转向量模型
rna2vec = BaseSeq2Vec(base_info.get_RNAf().shape[-1], 32, base_info.get_rna_s2i(), base_info.get_RNAf(), base_info.k)
dna2vec = BaseSeq2Vec(base_info.get_DNAf().shape[-1], 32, base_info.get_dna_s2i(), base_info.get_DNAf(), base_info.k)

# 开始与训练并保存
print('training rna2vec')
rna2vec.pretrain(rna_seq_list, 5)
rna2vec.save(f'{OUTPUT_PATH}rna2vec.m')

print('training dna2vec')
dna2vec.pretrain(dna_seq_list, 5)
dna2vec.save(f'{OUTPUT_PATH}dna2vec.m')
